import pywt
import numpy as np
import matplotlib.pyplot as plt

def wavelet_denoise(signal, wavelet='db2', levels=2, threshold_mode='soft'):
    """
    小波降噪：在时频域中雕刻信号的纯净形态
    
    这个函数体现了三个核心模式：
    1. 分解：将复杂信号拆解为简单组件
    2. 筛选：在每个尺度上识别并保留重要信息
    3. 重构：将纯净的组件重新组合
    """
    
    # 第一乐章：多尺度分解
    coeffs = pywt.wavedec(signal, wavelet, level=levels)
    
    # 第二乐章：智能阈值估计（使用贝叶斯风险准则）
    sigma = np.median(np.abs(coeffs[-1])) / 0.6745  # 噪声标准差估计
    threshold = sigma * np.sqrt(2 * np.log(len(signal)))  # 通用阈值
    
    # 第三乐章：分层净化
    coeffs_thresh = list(coeffs)
    coeffs_thresh[0] = coeffs[0]  # 保留低频近似（主旋律）
    
    for i in range(1, len(coeffs)):
        if threshold_mode == 'soft':
            coeffs_thresh[i] = pywt.threshold(coeffs[i], threshold, mode='soft')
        else:
            coeffs_thresh[i] = pywt.threshold(coeffs[i], threshold, mode='hard')
    
    # 第四乐章：完美重构
    denoised_signal = pywt.waverec(coeffs_thresh, wavelet)
    
    return denoised_signal, coeffs, coeffs_thresh, threshold

def _create_simple_signal(num_samples=1000, freq1=5, freq2=20, noise_level=0.3):
    """生成简单的双频正弦波+噪声信号"""
    t = np.linspace(0, 1, num_samples)
    clean = np.sin(2 * np.pi * freq1 * t) + 0.5 * np.sin(2 * np.pi * freq2 * t)
    noise = noise_level * np.random.randn(len(t))
    return t, clean + noise, clean

def _create_seismic_like_signal(
    num_samples=1000,          # 信号点数
    base_amplitude=0.0,        # 信号基线（中心值）
    wave_amplitude=0.1,        # 波动幅度（控制起伏程度）
    wave_freqs=(1, 3),         # 波动频率范围 (Hz)
    random_noise_amplitude=0.0 # 随机噪声幅度
):
    """
    生成简化版类地震信号，仅包含基础基线和少量低频正弦/余弦波动
    Returns:
        tuple: (t, noisy_signal, clean_signal)
    """
    t = np.arange(num_samples)
    
    # 纯净信号：仅基线 + 低频波动
    clean_signal = np.full(num_samples, base_amplitude)
    
    # 叠加1-2个低频正弦/余弦波
    noisy_signal = clean_signal
    for freq in np.linspace(wave_freqs[0], wave_freqs[1], 2):
        noisy_signal += wave_amplitude * np.sin(2 * np.pi * freq * t / num_samples)
        noisy_signal += wave_amplitude * 0.5 * np.cos(2 * np.pi * (freq + 1) * t / num_samples)
    
    return t, noisy_signal, clean_signal
# 演示：创造一个含噪信号并净化它
def demonstrate_denoising():
    # 创建测试信号：正弦波 + 噪声
    t, noisy_signal, clean_signal = _create_seismic_like_signal()
    # 应用小波降噪
    denoised, orig_coeffs, thresh_coeffs, threshold = wavelet_denoise(
        noisy_signal, wavelet='db2', levels=2, threshold_mode='hard'
    )
    
    # 可视化：展现降噪的魔法
    fig, axes = plt.subplots(3, 2, figsize=(15, 12))
    
    # 时域对比
    axes[0, 0].plot(t, clean_signal, 'g-', label='原始信号', linewidth=2)
    axes[0, 0].plot(t, noisy_signal, 'r-', alpha=0.7, label='含噪信号')
    axes[0, 0].set_title('信号对比：纯净 vs 污染')
    axes[0, 0].legend()
    axes[0, 0].grid(True)
    
    axes[0, 1].plot(t, clean_signal, 'g-', label='原始信号', linewidth=2)
    axes[0, 1].plot(t, denoised, 'b--', label='降噪信号', linewidth=2)
    axes[0, 1].set_title('降噪效果：重获新生')
    axes[0, 1].legend()
    axes[0, 1].grid(True)
    
    # 小波系数对比
    for i, (orig, thresh) in enumerate(zip(orig_coeffs[1:3], thresh_coeffs[1:3])):
        axes[1, i].stem(orig, linefmt='r-', markerfmt='ro', basefmt=' ', label='原始系数')
        axes[1, i].stem(thresh, linefmt='b-', markerfmt='bo', basefmt=' ', label='阈值化后')
        axes[1, i].axhline(y=threshold, color='k', linestyle='--', label=f'阈值={threshold:.3f}')
        axes[1, i].axhline(y=-threshold, color='k', linestyle='--')
        axes[1, i].set_title(f'第{i+1}层细节系数')
        axes[1, i].legend()
        axes[1, i].grid(True)
    
    # 频谱分析
    freqs = np.fft.fftfreq(len(t), t[1] - t[0])
    fft_clean = np.fft.fft(clean_signal)
    fft_noisy = np.fft.fft(noisy_signal)
    fft_denoised = np.fft.fft(denoised)
    
    axes[2, 0].semilogy(freqs[:len(freqs)//2], np.abs(fft_noisy[:len(freqs)//2]), 'r-', alpha=0.7, label='含噪信号')
    axes[2, 0].semilogy(freqs[:len(freqs)//2], np.abs(fft_clean[:len(freqs)//2]), 'g-', linewidth=2, label='原始信号')
    axes[2, 0].set_title('频域：噪声的指纹')
    axes[2, 0].set_xlabel('频率 (Hz)')
    axes[2, 0].set_ylabel('幅度')
    axes[2, 0].legend()
    axes[2, 0].grid(True)
    
    axes[2, 1].semilogy(freqs[:len(freqs)//2], np.abs(fft_denoised[:len(freqs)//2]), 'b-', linewidth=2, label='降噪信号')
    axes[2, 1].semilogy(freqs[:len(freqs)//2], np.abs(fft_clean[:len(freqs)//2]), 'g--', linewidth=2, label='原始信号')
    axes[2, 1].set_title('频域：重获纯净')
    axes[2, 1].set_xlabel('频率 (Hz)')
    axes[2, 1].set_ylabel('幅度')
    axes[2, 1].legend()
    axes[2, 1].grid(True)
    
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 更好的中文字体，可以显示“-”号
    plt.tight_layout()
    #plt.savefig('wavelet_denoising_analysis.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    # 性能指标
    mse_noisy = np.mean((noisy_signal - clean_signal)**2)
    mse_denoised = np.mean((denoised - clean_signal)**2)
    snr_improvement = 10 * np.log10(mse_noisy / mse_denoised)
    
    print(f"\n🎯 降噪性能分析:")
    print(f"   原始信噪比提升: {snr_improvement:.2f} dB")
    print(f"   均方误差减少: {(1 - mse_denoised/mse_noisy)*100:.1f}%")
    print(f"   使用阈值: {threshold:.4f}")
    
    return denoised, clean_signal, noisy_signal

if __name__ == "__main__":
    demonstrate_denoising()